#include <windows.h>
#include <iostream>
#include <string>
#include "detours.h"

using namespace std;

#ifdef _WIN64
EXTERN_C __declspec(dllexport) int AceInject(DWORD dwProcessId, BOOL wow64, LPVOID lpBaseAddr);
#else
EXTERN_C __declspec(dllexport) int AceInject(DWORD dwProcessId);
#endif

typedef struct 
{
    LPCSTR lpLibFileName;
    LPCSTR lpProcName;
    LPVOID pDetour;
    LPVOID* ppPointer;
} DETOUR_ENTITY, *PDETOUR_ENTITY, *LPDETOUR_ENTITY;
EXTERN_C __declspec(dllexport) int AceDetour(LPDETOUR_ENTITY lpEntityList, int size);


#define DLL_CORE_X86 "core-x86"
#define DLL_CORE_X64 "core-x64"


#ifdef _WIN64
int AceInject(DWORD dwProcessId, BOOL wow64, LPVOID lpBaseAddr)
#else
int AceInject(DWORD dwProcessId)
#endif
{
#ifdef _WIN64
    if(!dwProcessId || (wow64 && !lpBaseAddr))
    {
        SetLastError(ERROR_INVALID_PARAMETER);
        return FALSE;
    }
    if(!wow64)
    {
        lpBaseAddr = LoadLibraryA;
    }
    CHAR szPayloadX86[] = DLL_CORE_X86;
    CHAR szPayloadX64[] = DLL_CORE_X64;
    CHAR* szPayload = wow64 ? szPayloadX86 : szPayloadX64;
#else
    if(!dwProcessId)
    {
        SetLastError(ERROR_INVALID_PARAMETER);
        return FALSE;
    }
    LPVOID lpBaseAddr = LoadLibraryA;
    CHAR szPayload[] = DLL_CORE_X86;
#endif

    const DWORD dwMaxWaitTime = 5000;
    BOOL result = TRUE;
    HANDLE hProcess = NULL;
    SIZE_T dwSize;
    LPVOID lpStartAddress = NULL;
    LPVOID lpParameter = NULL;
    HANDLE hThread = NULL;
    DWORD dwResult;

    // open process
    hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, dwProcessId);
    if(!hProcess)
    {
        result = FALSE;
        goto exit;
    }

    // malloc on target
    dwSize = strlen(szPayload) * sizeof(CHAR);
    lpStartAddress = lpBaseAddr;
    lpParameter = VirtualAllocEx(hProcess, NULL, dwSize, MEM_COMMIT, PAGE_READWRITE);
    if(!lpParameter)
    {
        result = FALSE;
        goto exit;
    }

    // write target memory
    if(!WriteProcessMemory(hProcess, lpParameter, szPayload, dwSize, NULL))
    {
        result = FALSE;
        goto exit;
    }

    // create remote thread
    hThread = CreateRemoteThread(hProcess, NULL, NULL, (LPTHREAD_START_ROUTINE)lpStartAddress, lpParameter, NULL, NULL);    
    if(!hThread)
    {
        result = FALSE;
        goto exit;
    }

    // wait for ready
    dwResult = WaitForSingleObject(hThread, dwMaxWaitTime);
    switch(dwResult)
    {
    case WAIT_OBJECT_0:
        {
            // this might just have worked, pick up the result!
            result = TRUE;
        }
        break;
    case WAIT_TIMEOUT:
    case WAIT_FAILED:
    default:
        {
            result = FALSE;
            goto exit;
        }
    }

exit:
    if(lpParameter)
        VirtualFreeEx(hProcess, lpParameter, 0, MEM_RELEASE);
    if(hThread)
        CloseHandle(hThread);
    if(hProcess)
        CloseHandle(hProcess);
    return result;
}

int AceDetour(LPDETOUR_ENTITY lpEntityList, int size)
{
    if(!lpEntityList || size < 0)
    {
        SetLastError(ERROR_INVALID_PARAMETER);
        return FALSE;
    }

    // TODO:
    for (int i = 0; i < size; i++)
    {
        //cout << lpEntityList[i].lpLibFileName << endl;
        //cout << lpEntityList[i].lpProcName << endl;
    }

    return TRUE;
}

